#ifndef MULTIGRID_TWO_DIMENSION_HPP
#define MULTIGRID_TWO_DIMENSION_HPP

#include <iostream>
#include <cmath>
#include <vector>
#include <algorithm>
#include <fstream>
#include <Eigen/Sparse>
#include <Eigen/Dense>
#include <unsupported/Eigen/KroneckerProduct>
#include "Poisson.hpp"
#include "Function.hpp"
#include "json.hpp"
#include <chrono>

class Timer
{
public:
    Timer() : start_(std::chrono::high_resolution_clock::now()) {}

    void reset()
    {
        start_ = std::chrono::high_resolution_clock::now();
    }

    double elapsed() const
    {
        auto end = std::chrono::high_resolution_clock::now();
        return std::chrono::duration_cast<std::chrono::milliseconds>(end - start_).count();
    }

private:
    std::chrono::time_point<std::chrono::high_resolution_clock> start_;
};

struct BoundaryConditions_Two_Dimension
{
    enum Type
    {
        DIRICHLET,
        NEUMANN,
        Mixed
    };
    Type type;
    std::vector<int> flag;
    std::vector<double> one_point;
    // std::vector<double> value;
    std::shared_ptr<Function_Two_Dimension> u_1;
    std::shared_ptr<Function_Two_Dimension> u_2;
    std::shared_ptr<Function_Two_Dimension> u_3;
    std::shared_ptr<Function_Two_Dimension> u_4;
};

struct MGParams_Two_Dimension
{
    int grid_resolution;
    enum Restriction
    {
        full_weighting,
        injection
    };
    Restriction restriction;
    enum Interpolation
    {
        linear,
        quadratic
    };
    Interpolation interpolation;
    enum Type
    {
        V_Cycles,
        FMG
    };
    Type type;
    int max_iterations;
    int recursion_limit = -1;
    double relative_accuracy;
    std::shared_ptr<Function_Two_Dimension> initial_guess;
    int v1; // Forward Relax times
    int v2; // Backward Relax times
};

class Multigrid_Two_Dimension
{
private:
    MGParams_Two_Dimension MG;
    BoundaryConditions_Two_Dimension BC;
    std::shared_ptr<Function_Two_Dimension> source_term;
    double h;
    double h2;
    int n;
    int m;
    double omega = 0.75; // Jacobi权重系数
    Eigen::SparseMatrix<double> A;
    Eigen::VectorXd f;
    Eigen::VectorXd U;
    std::vector<int> ipiv;
    std::vector<std::vector<int>> index_map;

    Eigen::SparseMatrix<double> I_h_2h_1d(int N)
    {
        Eigen::SparseMatrix<double> X(N / 2 + 1, N + 1);
        int flag = 1;
        if (MG.restriction == MGParams_Two_Dimension::injection)
        {
            for (int i = 0; i < N / 2 + 1; i++)
            {
                X.insert(i, 2 * i) = 1;
            }
        }
        if (MG.restriction == MGParams_Two_Dimension::full_weighting)
        {
            if (BC.type == BoundaryConditions_Two_Dimension::DIRICHLET)
            {
                X.insert(0, 0) = 1;
                X.insert(N / 2, N) = 1;
                for (int i = 1; i < N / 2; i++)
                {
                    X.insert(i, flag++) = 0.25;
                    X.insert(i, flag++) = 0.5;
                    X.insert(i, flag) = 0.25;
                }
            }
            if (BC.type == BoundaryConditions_Two_Dimension::NEUMANN)
            {
                X.insert(0, 0) = 0.5;
                X.insert(0, 1) = 0.5;
                X.insert(N / 2, N - 1) = 0.5;
                X.insert(N / 2, N) = 0.5;
                for (int i = 1; i < N / 2; i++)
                {
                    X.insert(i, flag++) = 0.25;
                    X.insert(i, flag++) = 0.5;
                    X.insert(i, flag) = 0.25;
                }
            }
        }
        return X;
    }

    Eigen::SparseMatrix<double> I_2h_h_1d(int N)
    {
        Eigen::SparseMatrix<double> X(N + 1, N / 2 + 1);
        if (MG.interpolation == MGParams_Two_Dimension::quadratic)
        {
            X = I_h_2h_1d(N).transpose();
        }
        if (MG.interpolation == MGParams_Two_Dimension::linear)
        {
            int flag = 0;
            for (int i = 0; i < N / 2; i++)
            {
                X.insert(flag++, i) = 0.5;
                X.insert(flag++, i) = 1.0;
                X.insert(flag, i) = 0.5;
            }
        }
        return X;
    }

    Eigen::SparseMatrix<double> I_h_2h(int N)
    {
        Eigen::SparseMatrix<double> X((N / 2 + 1) * (N / 2 + 1), (N + 1) * (N + 1));
        if (MG.restriction == MGParams_Two_Dimension::injection)
        {
            if (BC.type == BoundaryConditions_Two_Dimension::DIRICHLET)
            {
                X = Eigen::kroneckerProduct(I_h_2h_1d(N), I_h_2h_1d(N));
            }
            if (BC.type == BoundaryConditions_Two_Dimension::NEUMANN)
            {
                X = Eigen::kroneckerProduct(I_h_2h_1d(N), I_h_2h_1d(N));
            }
            if (BC.type == BoundaryConditions_Two_Dimension::Mixed)
            {
                for (int j = 0; j < N / 2 + 1; j++)
                {
                    for (int i = 0; i < N / 2 + 1; i++)
                    {
                        int coarse_idx = i + j * (N / 2 + 1);
                        int fine_idx = (2 * i) + (2 * j) * (N + 1);
                        X.insert(coarse_idx, fine_idx) = 1.0;
                    }
                }
            }
            else
            {
                X = Eigen::kroneckerProduct(I_h_2h_1d(N), I_h_2h_1d(N));
            }
        }
        if (MG.restriction == MGParams_Two_Dimension::full_weighting)
        {
            if (BC.type == BoundaryConditions_Two_Dimension::DIRICHLET)
            {
                X = Eigen::kroneckerProduct(I_h_2h_1d(N), I_h_2h_1d(N));
            }
            if (BC.type == BoundaryConditions_Two_Dimension::NEUMANN)
            {
                X = Eigen::kroneckerProduct(I_h_2h_1d(N), I_h_2h_1d(N));
            }
            if (BC.type == BoundaryConditions_Two_Dimension::Mixed)
            {
                for (int j = 0; j < N / 2 + 1; j++)
                {
                    for (int i = 0; i < N / 2 + 1; i++)
                    {
                        int coarse_idx = i + j * (N / 2 + 1);
                        int fi = 2 * i;
                        int fj = 2 * j;
                        int center = fi + fj * (N + 1);

                        // Different weights for different boundary types
                        bool is_dirichlet = false;
                        bool is_neumann = false;

                        // Check boundary types
                        if (BC.flag[0] == 0 && j == 0)
                            is_dirichlet = true; // Bottom boundary
                        if (BC.flag[1] == 0 && i == N)
                            is_dirichlet = true; // Right boundary
                        if (BC.flag[2] == 0 && j == N)
                            is_dirichlet = true; // Top boundary
                        if (BC.flag[3] == 0 && i == 0)
                            is_dirichlet = true; // Left boundary

                        if (BC.flag[0] == 1 && j == 0)
                            is_neumann = true; // Bottom boundary
                        if (BC.flag[1] == 1 && i == N)
                            is_neumann = true; // Right boundary
                        if (BC.flag[2] == 1 && j == N)
                            is_neumann = true; // Top boundary
                        if (BC.flag[3] == 1 && i == 0)
                            is_neumann = true; // Left boundary

                        if (is_dirichlet)
                        {
                            // Dirichlet boundary - use injection
                            X.insert(coarse_idx, center) = 1.0;
                        }
                        else if (is_neumann)
                        {
                            // Neumann boundary - special treatment
                            if (i == 0 && j == 0)
                            {
                                // Bottom-left corner
                                X.insert(coarse_idx, center) = 0.5;
                                X.insert(coarse_idx, center + 1) = 0.25;
                                X.insert(coarse_idx, center + (N + 1)) = 0.25;
                            }
                            else if (i == N / 2 && j == 0)
                            {
                                // Bottom-right corner
                                X.insert(coarse_idx, center) = 0.5;
                                X.insert(coarse_idx, center - 1) = 0.25;
                                X.insert(coarse_idx, center + (N + 1)) = 0.25;
                            }
                            else if (i == 0 && j == N / 2)
                            {
                                // Top-left corner
                                X.insert(coarse_idx, center) = 0.5;
                                X.insert(coarse_idx, center + 1) = 0.25;
                                X.insert(coarse_idx, center - (N + 1)) = 0.25;
                            }
                            else if (i == N / 2 && j == N / 2)
                            {
                                // Top-right corner
                                X.insert(coarse_idx, center) = 0.5;
                                X.insert(coarse_idx, center - 1) = 0.25;
                                X.insert(coarse_idx, center - (N + 1)) = 0.25;
                            }
                            else if (j == 0)
                            {
                                // Bottom boundary (not corner)
                                X.insert(coarse_idx, center) = 0.5;
                                X.insert(coarse_idx, center - 1) = 0.125;
                                X.insert(coarse_idx, center + 1) = 0.125;
                                X.insert(coarse_idx, center + (N + 1)) = 0.25;
                            }
                            else if (i == N / 2)
                            {
                                // Right boundary (not corner)
                                X.insert(coarse_idx, center) = 0.5;
                                X.insert(coarse_idx, center - (N + 1)) = 0.125;
                                X.insert(coarse_idx, center + (N + 1)) = 0.125;
                                X.insert(coarse_idx, center - 1) = 0.25;
                            }
                            else if (j == N / 2)
                            {
                                // Top boundary (not corner)
                                X.insert(coarse_idx, center) = 0.5;
                                X.insert(coarse_idx, center - 1) = 0.125;
                                X.insert(coarse_idx, center + 1) = 0.125;
                                X.insert(coarse_idx, center - (N + 1)) = 0.25;
                            }
                            else if (i == 0)
                            {
                                // Left boundary (not corner)
                                X.insert(coarse_idx, center) = 0.5;
                                X.insert(coarse_idx, center - (N + 1)) = 0.125;
                                X.insert(coarse_idx, center + (N + 1)) = 0.125;
                                X.insert(coarse_idx, center + 1) = 0.25;
                            }
                        }
                        else
                        {
                            // Interior point - standard full weighting
                            X.insert(coarse_idx, center) = 0.25;
                            if (center >= 1)
                                X.insert(coarse_idx, center - 1) = 0.125;
                            if (center <= (N + 1) * (N + 1) - 2)
                                X.insert(coarse_idx, center + 1) = 0.125;
                            if (center >= N + 1)
                                X.insert(coarse_idx, center - (N + 1)) = 0.125;
                            if (center <= (N + 1) * (N + 1) - N - 2)
                                X.insert(coarse_idx, center + (N + 1)) = 0.125;
                            if (center >= N + 2)
                                X.insert(coarse_idx, center - (N + 1) - 1) = 0.0625;
                            if (center >= N)
                                X.insert(coarse_idx, center - (N + 1) + 1) = 0.0625;
                            if (center <= (N + 1) * (N + 1) - N - 1)
                                X.insert(coarse_idx, center + (N + 1) - 1) = 0.0625;
                            if (center <= (N + 1) * (N + 1) - N - 3)
                                X.insert(coarse_idx, center + (N + 1) + 1) = 0.0625;
                        }
                    }
                }
            }
            else
            {
                X = Eigen::kroneckerProduct(I_h_2h_1d(N), I_h_2h_1d(N));
            }
        }
        return X;
    }

    Eigen::SparseMatrix<double> I_2h_h(int N)
    {
        Eigen::SparseMatrix<double> X((N + 1) * (N + 1), (N / 2 + 1) * (N / 2 + 1));
        if (MG.interpolation == MGParams_Two_Dimension::quadratic)
        {
            int flag = 1;
            Eigen::SparseMatrix<double> X1((N / 2 + 1), (N + 1));
            if (BC.type == BoundaryConditions_Two_Dimension::DIRICHLET)
            {
                X1.insert(0, 0) = 1;
                X1.insert(N / 2, N) = 1;
                for (int i = 1; i < N / 2; i++)
                {
                    X1.insert(i, flag++) = 0.25;
                    X1.insert(i, flag++) = 0.5;
                    X1.insert(i, flag) = 0.25;
                }
                X = Eigen::kroneckerProduct(X1, X1).transpose();
            }
            if (BC.type == BoundaryConditions_Two_Dimension::NEUMANN)
            {
                X1.insert(0, 0) = 0.5;
                X1.insert(0, 1) = 0.5;
                X1.insert(N / 2, N - 1) = 0.5;
                X1.insert(N / 2, N) = 0.5;
                for (int i = 1; i < N / 2; i++)
                {
                    X1.insert(i, flag++) = 0.25;
                    X1.insert(i, flag++) = 0.5;
                    X1.insert(i, flag) = 0.25;
                }
                X = Eigen::kroneckerProduct(X1, X1).transpose();
            }
            if (BC.type == BoundaryConditions_Two_Dimension::Mixed)
            {
                for (int j = 0; j < N / 2 + 1; j++)
                {
                    for (int i = 0; i < N / 2 + 1; i++)
                    {
                        int coarse_idx = i + j * (N / 2 + 1);
                        int fi = 2 * i;
                        int fj = 2 * j;
                        int fine_center = fi + fj * (N + 1);

                        bool is_dirichlet = false;
                        bool is_neumann = false;

                        // Check boundary types
                        if (BC.flag[0] == 0 && j == 0)
                            is_dirichlet = true; // Bottom boundary
                        if (BC.flag[1] == 0 && i == N / 2)
                            is_dirichlet = true; // Right boundary
                        if (BC.flag[2] == 0 && j == N / 2)
                            is_dirichlet = true; // Top boundary
                        if (BC.flag[3] == 0 && i == 0)
                            is_dirichlet = true; // Left boundary

                        if (BC.flag[0] == 1 && j == 0)
                            is_neumann = true; // Bottom boundary
                        if (BC.flag[1] == 1 && i == N / 2)
                            is_neumann = true; // Right boundary
                        if (BC.flag[2] == 1 && j == N / 2)
                            is_neumann = true; // Top boundary
                        if (BC.flag[3] == 1 && i == 0)
                            is_neumann = true; // Left boundary

                        if (is_dirichlet)
                        {
                            // Dirichlet boundary - direct injection
                            X.insert(fine_center, coarse_idx) = 1.0;
                        }
                        else if (is_neumann)
                        {
                            // Neumann boundary - special treatment
                            if (i == 0 && j == 0)
                            {
                                // Bottom-left corner
                                X.insert(fine_center, coarse_idx) = 1.0;
                                X.insert(fine_center + 1, coarse_idx) = 0.5;
                                X.insert(fine_center + (N + 1), coarse_idx) = 0.5;
                            }
                            else if (i == N / 2 && j == 0)
                            {
                                // Bottom-right corner
                                X.insert(fine_center, coarse_idx) = 1.0;
                                X.insert(fine_center - 1, coarse_idx) = 0.5;
                                X.insert(fine_center + (N + 1), coarse_idx) = 0.5;
                            }
                            else if (i == 0 && j == N / 2)
                            {
                                // Top-left corner
                                X.insert(fine_center, coarse_idx) = 1.0;
                                X.insert(fine_center + 1, coarse_idx) = 0.5;
                                X.insert(fine_center - (N + 1), coarse_idx) = 0.5;
                            }
                            else if (i == N / 2 && j == N / 2)
                            {
                                // Top-right corner
                                X.insert(fine_center, coarse_idx) = 1.0;
                                X.insert(fine_center - 1, coarse_idx) = 0.5;
                                X.insert(fine_center - (N + 1), coarse_idx) = 0.5;
                            }
                            else if (j == 0)
                            {
                                // Bottom boundary (not corner)
                                X.insert(fine_center, coarse_idx) = 1.0;
                                X.insert(fine_center - 1, coarse_idx) = 0.5;
                                X.insert(fine_center + 1, coarse_idx) = 0.5;
                                X.insert(fine_center + (N + 1), coarse_idx) = 0.5;
                            }
                            else if (i == N / 2)
                            {
                                // Right boundary (not corner)
                                X.insert(fine_center, coarse_idx) = 1.0;
                                X.insert(fine_center - (N + 1), coarse_idx) = 0.5;
                                X.insert(fine_center + (N + 1), coarse_idx) = 0.5;
                                X.insert(fine_center - 1, coarse_idx) = 0.5;
                            }
                            else if (j == N / 2)
                            {
                                // Top boundary (not corner)
                                X.insert(fine_center, coarse_idx) = 1.0;
                                X.insert(fine_center - 1, coarse_idx) = 0.5;
                                X.insert(fine_center + 1, coarse_idx) = 0.5;
                                X.insert(fine_center - (N + 1), coarse_idx) = 0.5;
                            }
                            else if (i == 0)
                            {
                                // Left boundary (not corner)
                                X.insert(fine_center, coarse_idx) = 1.0;
                                X.insert(fine_center - (N + 1), coarse_idx) = 0.5;
                                X.insert(fine_center + (N + 1), coarse_idx) = 0.5;
                                X.insert(fine_center + 1, coarse_idx) = 0.5;
                            }
                        }
                        else
                        {
                            // Interior point - standard quadratic interpolation
                            X.insert(fine_center, coarse_idx) = 1.0;
                            if (fi > 0)
                                X.insert(fine_center - 1, coarse_idx) = 0.5;
                            if (fi < N)
                                X.insert(fine_center + 1, coarse_idx) = 0.5;
                            if (fj > 0)
                                X.insert(fine_center - (N + 1), coarse_idx) = 0.5;
                            if (fj < N)
                                X.insert(fine_center + (N + 1), coarse_idx) = 0.5;
                            if (fi > 1)
                                X.insert(fine_center - 2, coarse_idx) = 0.25;
                            if (fi < N - 1)
                                X.insert(fine_center + 2, coarse_idx) = 0.25;
                            if (fj > 1)
                                X.insert(fine_center - 2 * (N + 1), coarse_idx) = 0.25;
                            if (fj < N - 1)
                                X.insert(fine_center + 2 * (N + 1), coarse_idx) = 0.25;
                        }
                    }
                }
            }
            else
            {
                X = Eigen::kroneckerProduct(I_h_2h_1d(N), I_h_2h_1d(N)).transpose();
            }
        }

        if (MG.interpolation == MGParams_Two_Dimension::linear)
        {
            // Bilinear interpolation
            for (int j = 0; j < N / 2 + 1; j++)
            {
                for (int i = 0; i < N / 2 + 1; i++)
                {
                    int coarse_idx = i + j * (N / 2 + 1);

                    // Corresponding fine grid positions
                    int fi = 2 * i;
                    int fj = 2 * j;
                    int fine_center = fi + fj * (N + 1);

                    // Center point (coincides with coarse point)
                    X.insert(fine_center, coarse_idx) = 1.0;

                    // Horizontal neighbors
                    if (fi > 0)
                        X.insert(fine_center - 1, coarse_idx) = 0.5;
                    if (fi < N)
                        X.insert(fine_center + 1, coarse_idx) = 0.5;

                    // Vertical neighbors
                    if (fj > 0)
                        X.insert(fine_center - (N + 1), coarse_idx) = 0.5;
                    if (fj < N)
                        X.insert(fine_center + (N + 1), coarse_idx) = 0.5;

                    // Corner neighbors (bilinear interpolation)
                    if (fi > 0 && fj > 0)
                        X.insert(fine_center - (N + 1) - 1, coarse_idx) = 0.25;
                    if (fi < N && fj > 0)
                        X.insert(fine_center - (N + 1) + 1, coarse_idx) = 0.25;
                    if (fi > 0 && fj < N)
                        X.insert(fine_center + (N + 1) - 1, coarse_idx) = 0.25;
                    if (fi < N && fj < N)
                        X.insert(fine_center + (N + 1) + 1, coarse_idx) = 0.25;
                }
            }
        }
        return X;
    }


    void initialization()
    {
        int total_points = (m + 2) * (m + 2);
        n = m + 1;
        A.resize(total_points, total_points);
        A.reserve(Eigen::VectorXi::Constant(total_points, 5));
        U.resize(total_points);
        f.resize(total_points);
        ipiv.resize(total_points, 0);
        if (MG.type == MGParams_Two_Dimension::V_Cycles)
        {
            for (int i = 0; i < m + 2; ++i)
            {
                for (int j = 0; j < m + 2; ++j)
                    U[i] = MG.initial_guess->operator()((i)*h, j * h);
            }
        }

        if (BC.type == BoundaryConditions_Two_Dimension::DIRICHLET)
        {
            for (int i = 0; i < total_points; ++i)
            {
                int w = i % (m + 2);
                if (i <= m + 1)
                {
                    A.insert(i, i) = 4 / h2;
                }
                else if (i >= total_points - m - 2)
                {
                    A.insert(i, i) = 4 / h2;
                }
                else if (w == 0 || w == m + 1)
                {
                    A.insert(i, i) = 4 / h2;
                }
                else
                {
                    A.insert(i, i) = 4.0 / h2;
                    A.insert(i, i - 1) = A.insert(i, i + 1) = -1.0 / h2;
                    A.insert(i, i - m - 2) = A.insert(i, i + m + 2) = -1.0 / h2;
                }
                for (int j = 0; j < m + 2; ++j)
                {
                    for (int i = 0; i < m + 2; ++i)
                    {
                        if (j == 0)
                        {
                            f[i + j * (m + 2)] = 4 * BC.u_1->operator()(i * h, j * h) / h2;
                        }
                        else if (i == m + 1)
                        {
                            f[i + j * (m + 2)] = 4 * BC.u_2->operator()(i * h, j * h) / h2;
                        }
                        else if (j == m + 1)
                        {
                            f[i + j * (m + 2)] = 4 * BC.u_3->operator()(i * h, j * h) / h2;
                        }
                        else if (i == 0)
                        {
                            f[i + j * (m + 2)] = 4 * BC.u_4->operator()(i * h, j * h) / h2;
                        }
                        else
                        {
                            f[i + j * (m + 2)] = (*source_term)(i * h, j * h);
                        }
                    }
                }
            }
        }

        if (BC.type == BoundaryConditions_Two_Dimension::NEUMANN)
        {
            if (BC.one_point.size() != 3)
            {
                throw std::runtime_error("BoundaryConditions_Two_Dimension one_point vector must be 3!");
            }
            int c = 0;
            for (int j = 0; j < m + 2; ++j)
            {
                for (int i = 0; i < m + 2; ++i)
                {
                    c = i + j * (m + 2);
                    if (i == 0 && j == 0)
                    {
                        A.insert(c, c) = 4.0 / h2;
                        A.insert(c, c + 1) = -2.0 / h2;
                        A.insert(c, c + m + 2) = -2.0 / h2;
                        f[i + j * (m + 2)] = -(2.0 * BC.u_4->operator()(i * h, j * h)) / h - (2.0 * BC.u_1->operator()(i * h, j * h)) / h + (*source_term)(i * h, j * h);
                    }
                    else if (j == 0 && i == m + 1)
                    {
                        A.insert(c, c) = 4.0 / h2;
                        A.insert(c, c - 1) = -2.0 / h2;
                        A.insert(c, c + m + 2) = -2.0 / h2;
                        f[i + j * (m + 2)] = (2.0 * BC.u_2->operator()(i * h, j * h)) / h - (2.0 * BC.u_1->operator()(i * h, j * h)) / h + (*source_term)(i * h, j * h);
                    }
                    else if (j == m + 1 && i == 0)
                    {
                        A.insert(c, c) = 4.0 / h2;
                        A.insert(c, c + 1) = -2.0 / h2;
                        A.insert(c, c - m - 2) = -2.0 / h2;
                        f[i + j * (m + 2)] = -(2.0 * BC.u_4->operator()(i * h, j * h)) / h + (2.0 * BC.u_3->operator()(i * h, j * h)) / h + (*source_term)(i * h, j * h);
                    }
                    else if (j == m + 1 && i == m + 1)
                    {
                        A.insert(c, c) = 4.0 / h2;
                        A.insert(c, c - 1) = -2.0 / h2;
                        A.insert(c, c - m - 2) = -2.0 / h2;
                        f[i + j * (m + 2)] = (2.0 * BC.u_3->operator()(i * h, j * h)) / h + (2.0 * BC.u_2->operator()(i * h, j * h)) / h + (*source_term)(i * h, j * h);
                    }
                    else if (j == 0 && i >= 1 && i <= m)
                    {
                        A.insert(c, c) = 4.0 / h2;
                        A.insert(c, c - 1) = -1.0 / (h2);
                        A.insert(c, c + 1) = -1.0 / (h2);
                        A.insert(c, c + m + 2) = -2.0 / (h2);
                        f[i + j * (m + 2)] = -(2.0 * BC.u_1->operator()(i * h, j * h)) / h + (*source_term)(i * h, j * h);
                    }
                    else if (i == m + 1 && j >= 1 && j <= m)
                    {
                        A.insert(c, c) = 4.0 / h2;
                        A.insert(c, c - 1) = -2.0 / h2;
                        A.insert(c, c + m + 2) = -1.0 / (h2);
                        A.insert(c, c - m - 2) = -1.0 / (h2);
                        f[i + j * (m + 2)] = (2.0 * BC.u_2->operator()(i * h, j * h)) / h + (*source_term)(i * h, j * h);
                    }
                    else if (j == m + 1 && i >= 1 && i <= m)
                    {
                        A.insert(c, c) = 4.0 / h2;
                        A.insert(c, c + 1) = -1.0 / (h2);
                        A.insert(c, c - 1) = -1.0 / (h2);
                        A.insert(c, c - m - 2) = -2.0 / (h2);
                        f[i + j * (m + 2)] = (2.0 * BC.u_3->operator()(i * h, j * h)) / h + (*source_term)(i * h, j * h);
                    }
                    else if (i == 0 && j >= 1 && j <= m)
                    {
                        A.insert(c, c) = 4.0 / h2;
                        A.insert(c, c + 1) = -2.0 / (h2);
                        A.insert(c, c + m + 2) = -1.0 / (h2);
                        A.insert(c, c - m - 2) = -1.0 / (h2);
                        f[i + j * (m + 2)] = -(2.0 * BC.u_4->operator()(i * h, j * h)) / h + (*source_term)(i * h, j * h);
                    }
                    else
                    {
                        A.insert(c, c) = 4.0 / h2;
                        A.insert(c, c + 1) = A.insert(c, c - 1) = -1.0 / h2;
                        A.insert(c, c + m + 2) = A.insert(c, c - m - 2) = -1.0 / h2;
                        f[i + j * (m + 2)] = (*source_term)(i * h, j * h);
                    }
                    if (i == int(BC.one_point[0] / h) && j == int(BC.one_point[1] / h))
                    {
                        A.coeffRef(c, c) += 4.0 / h2;
                        f[i + j * (m + 2)] += (4.0 * BC.one_point[2]) / h2;
                    }
                }
            }
        }

        if (BC.type == BoundaryConditions_Two_Dimension::Mixed)
        {
            if (BC.flag[0] == BC.flag[1] && BC.flag[0] == BC.flag[2] && BC.flag[0] == BC.flag[3])
            {
                std::cout << "Invalid Input for Mixed Boundary Conditions Two Dimension!!!" << std::endl;
            }
            for (int j = 0; j < m + 2; ++j)
            {
                for (int i = 0; i < m + 2; ++i)
                {
                    int c = i + j * (m + 2);
                    if (j == 0 && i == 0)
                    {
                        if (BC.flag[0] == 0)
                        {
                            A.insert(c, c) = 4.0 / h2;
                            f[c] = (4.0 * BC.u_1->operator()(i * h, j * h)) / h2;
                        }
                        else if (BC.flag[3] == 0)
                        {
                            A.insert(c, c) = 4.0 / h2;
                            f[c] = (4.0 * BC.u_4->operator()(i * h, j * h)) / h2;
                        }
                        else
                        {
                            A.insert(c, c) = 4.0 / h2;
                            A.insert(c, c + 1) = -2.0 / h2;
                            A.insert(c, c + m + 2) = -2.0 / h2;
                            f[c] = (*source_term)(i * h, j * h) - 2.0 * BC.u_4->operator()(i * h, j * h) / h - 2.0 * BC.u_1->operator()(i * h, j * h) / h;
                        }
                    }
                    else if (i == m + 1 && j == 0)
                    {
                        if (BC.flag[0] == 0)
                        {
                            A.insert(c, c) = 4.0 / h2;
                            f[c] = 4.0 * BC.u_1->operator()(i * h, j * h) / h2;
                        }
                        else if (BC.flag[1] == 0)
                        {
                            A.insert(c, c) = 4.0 / h2;
                            f[c] = 4.0 * BC.u_2->operator()(i * h, j * h) / h2;
                        }
                        else
                        {
                            A.insert(c, c) = 4.0 / h2;
                            A.insert(c, c - 1) = -2.0 / h2;
                            A.insert(c, c + m + 2) = -2.0 / h2;
                            f[c] = (*source_term)(i * h, j * h) + 2.0 * BC.u_2->operator()(i * h, j * h) / h - 2.0 * BC.u_1->operator()(i * h, j * h) / h;
                        }
                    }
                    else if (j == m + 1 && i == 0)
                    {
                        if (BC.flag[3] == 0)
                        {
                            A.insert(c, c) = 4.0 / h2;
                            f[c] = 4.0 * BC.u_4->operator()(i * h, j * h) / h2;
                        }
                        else if (BC.flag[2] == 0)
                        {
                            A.insert(c, c) = 4.0 / h2;
                            f[c] = 4.0 * BC.u_3->operator()(i * h, j * h) / h2;
                        }
                        else
                        {
                            A.insert(c, c) = 4.0 / h2;
                            A.insert(c, c + 1) = -2.0 / h2;
                            A.insert(c, c - m - 2) = -2.0 / h2;
                            f[c] = (*source_term)(i * h, j * h) - 2.0 * BC.u_4->operator()(i * h, j * h) / h + 2.0 * BC.u_3->operator()(i * h, j * h) / h;
                        }
                    }
                    else if (i == m + 1 && j == m + 1)
                    {
                        if (BC.flag[1] == 0)
                        {
                            A.insert(c, c) = 4.0 / h2;
                            f[c] = 4.0 * BC.u_2->operator()(i * h, j * h) / h2;
                        }
                        else if (BC.flag[2] == 0)
                        {
                            A.insert(c, c) = 4.0 / h2;
                            f[c] = 4.0 * BC.u_3->operator()(i * h, j * h) / h2;
                        }
                        else
                        {
                            A.insert(c, c) = 4.0 / h2;
                            A.insert(c, c - 1) = -2.0 / h2;
                            A.insert(c, c - m - 2) = -2.0 / h2;
                            f[i + j * (m + 2)] = (*source_term)(i * h, j * h) + 2.0 * BC.u_2->operator()(i * h, j * h) / h + 2.0 * BC.u_3->operator()(i * h, j * h) / h;
                        }
                    }
                    else if (i >= 1 && i <= m && j == 0)
                    {
                        if (BC.flag[0] == 0)
                        {
                            A.insert(c, c) = 4.0 / h2;
                            f[i + j * (m + 2)] = 4.0 * BC.u_1->operator()(i * h, j * h) / h2;
                        }
                        else
                        {
                            A.insert(c, c) = 4.0 / h2;
                            A.insert(c, c + 1) = -1.0 / h2;
                            A.insert(c, c - 1) = -1.0 / h2;
                            A.insert(c, c + m + 2) = -2.0 / h2;
                            f[i + j * (m + 2)] = (*source_term)(i * h, j * h) - 2.0 * BC.u_1->operator()(i * h, j * h) / h;
                        }
                    }
                    else if (i >= 1 && i <= m && j == m + 1)
                    {
                        if (BC.flag[2] == 0)
                        {
                            A.insert(c, c) = 4.0 / h2;
                            f[i + j * (m + 2)] = 4.0 * BC.u_3->operator()(i * h, j * h) / h2;
                        }
                        else
                        {
                            A.insert(c, c) = 4.0 / h2;
                            A.insert(c, c + 1) = -1.0 / h2;
                            A.insert(c, c - 1) = -1.0 / h2;
                            A.insert(c, c - m - 2) = -2.0 / h2;
                            f[i + j * (m + 2)] = (*source_term)(i * h, j * h) + 2.0 * BC.u_3->operator()(i * h, j * h) / h;
                        }
                    }
                    else if (j >= 1 && j <= m && i == 0)
                    {
                        if (BC.flag[3] == 0)
                        {
                            A.insert(c, c) = 4.0 / h2;
                            f[i + j * (m + 2)] = 4.0 * BC.u_4->operator()(i * h, j * h) / h2;
                        }
                        else
                        {
                            A.insert(c, c) = 4.0 / h2;
                            A.insert(c, c + m + 2) = -1.0 / h2;
                            A.insert(c, c - m - 2) = -1.0 / h2;
                            A.insert(c, c + 1) = -2.0 / h2;
                            f[i + j * (m + 2)] = (*source_term)(i * h, j * h) - 2.0 * BC.u_4->operator()(i * h, j * h) / h;
                        }
                    }
                    else if (j >= 1 && j <= m && i == m + 1)
                    {
                        if (BC.flag[1] == 0)
                        {
                            A.insert(c, c) = 4.0 / h2;
                            f[i + j * (m + 2)] = 4.0 * BC.u_2->operator()(i * h, j * h) / h2;
                        }
                        else
                        {
                            A.insert(c, c) = 4.0 / h2;
                            A.insert(c, c + m + 2) = -1.0 / h2;
                            A.insert(c, c - m - 2) = -1.0 / h2;
                            A.insert(c, c - 1) = -2.0 / h2;
                            f[i + j * (m + 2)] = (*source_term)(i * h, j * h) + 2.0 * BC.u_2->operator()(i * h, j * h) / h;
                        }
                    }
                    else
                    {
                        A.insert(c, c) = 4.0 / h2;
                        A.insert(c, c + m + 2) = -1.0 / h2;
                        A.insert(c, c - m - 2) = -1.0 / h2;
                        A.insert(c, c - 1) = -1.0 / h2;
                        A.insert(c, c + 1) = -1.0 / h2;
                        f[i + j * (m + 2)] = (*source_term)(i * h, j * h);
                    }
                }
            }
        }
        return;
    }

    // 加权 Jacobi 迭代
    void weightedJacobi(const Eigen::SparseMatrix<double> &A, Eigen::VectorXd &x, const Eigen::VectorXd &b, int v)
    {
        if (A.rows() != A.cols() || A.rows() != b.size() || b.size() != x.size())
        {
            throw std::runtime_error("Matrix and vector dimensions mismatch");
        }

        // Extract diagonal and compute its inverse
        Eigen::VectorXd D_inv = A.diagonal().cwiseInverse();

        // Create matrix R = L + U (off-diagonal parts)
        Eigen::SparseMatrix<double> R = A;
        R.diagonal().setZero();

        for (int k = 0; k < v; ++k)
        {
            x = (1.0 - omega) * x + omega * D_inv.cwiseProduct(b - R * x);
        }
        return;
    }

public:
    Multigrid_Two_Dimension(const MGParams_Two_Dimension &mg, const BoundaryConditions_Two_Dimension &bc, std::shared_ptr<Function_Two_Dimension> f)
        : MG(mg), BC(bc), source_term(f)
    {
        h = 1.0 / (MG.grid_resolution + 1);
        h2 = h * h;
        m = MG.grid_resolution;
        // std::cout << "here1" << std::endl;
        initialization();
        // std::cout << "here2" << std::endl;
        // std::cout << I_h_2h(4) << std::endl;
        // std::cout << "here3" << std::endl;
        // std::cout << I_2h_h(4) << std::endl;
        // std::cout << A << std::endl;
        // std::cout << f << std::endl;
    }

    Eigen::VectorXd V_cycles(Eigen::SparseMatrix<double> A1, Eigen::VectorXd v1, Eigen::VectorXd f1, int count_recursions)
    {

        const auto &I_h_2h_matrix = I_h_2h(std::sqrt(f1.size()) - 1);
        const auto &I_2h_h_matrix = I_2h_h(std::sqrt(f1.size()) - 1);
        // std::cout << f1.size() << std::endl;
        // std::cout << "A1:" << std::endl;
        // std::cout << A1 << std::endl;
        // std::cout << "I_2h_h:" << std::endl;
        // std::cout << I_2h_h_matrix << std::endl;
        // std::cout << ":::::::" << std::endl;
        weightedJacobi(A1, v1, f1, MG.v1);
        if (v1.size() <= 9)
        { // 最粗网格终止
            Eigen::SparseLU<Eigen::SparseMatrix<double>> solver;
            solver.compute(A1);
            v1 = solver.solve(f1);
            weightedJacobi(A1, v1, f1, MG.v2);
            return v1;
        }
        if (count_recursions == MG.recursion_limit)
        {
            Eigen::SparseLU<Eigen::SparseMatrix<double>> solver;
            solver.compute(A1);
            v1 = solver.solve(f1);
            weightedJacobi(A1, v1, f1, MG.v2);
            return v1;
        }
        Eigen::VectorXd f2 = I_h_2h_matrix * (f1 - A1 * v1);
        Eigen::VectorXd v2 = Eigen::VectorXd::Zero(f2.size());
        Eigen::SparseMatrix<double> A2 = I_h_2h_matrix * A1 * I_2h_h_matrix;
        v2 = V_cycles(A2, v2, f2, count_recursions + 1);
        v1 += I_2h_h_matrix * v2;
        weightedJacobi(A1, v1, f1, MG.v2);
        return v1;
    }

    Eigen::VectorXd FMG(Eigen::SparseMatrix<double> &A1, Eigen::VectorXd &f1, int count_recursions)
    {
        Eigen::VectorXd v1 = Eigen::VectorXd::Zero(f1.size());
        const auto &I_h_2h_matrix = I_h_2h(std::sqrt(f1.size()) - 1);
        const auto &I_2h_h_matrix = I_2h_h(std::sqrt(f1.size()) - 1);
        // std::cout << "count_recursions==" << count_recursions << std::endl;
        // std::cout << "A1:" << std::endl;
        // std::cout << A1 << std::endl;
        // std::cout << "I_h_2h:" << std::endl;
        // std::cout << I_h_2h_matrix << std::endl;
        // std::cout << ":::::::" << std::endl;
        if (std::sqrt(f1.size()) <= 9)
        { // 最粗网格终止
            v1 = V_cycles(A1, v1, f1, count_recursions);
            return v1;
        }
        if (count_recursions == MG.recursion_limit)
        {
            v1 = V_cycles(A1, v1, f1, count_recursions);
            return v1;
        }
        Eigen::VectorXd f2 = I_h_2h_matrix * (f1);
        Eigen::SparseMatrix<double> A2 = I_h_2h_matrix * A1 * I_2h_h_matrix;
        Eigen::VectorXd v2 = Eigen::VectorXd::Zero(f2.size());
        v2 = FMG(A2, f2, count_recursions + 1);
        v1 = I_2h_h_matrix * v2;
        v1 = V_cycles(A1, v1, f1, count_recursions);
        return v1;
    }


    void solve()
    {
        if (MG.type == MGParams_Two_Dimension::V_Cycles)
        {
            Timer timer;
            double temp = (f - A * U).norm() / f.norm();
            double relative_residual = (f - A * U).norm() / f.norm();
            for (int i = 0; i < MG.max_iterations; i++)
            {
                U = V_cycles(A, U, f, 0);
                relative_residual = (f - A * U).norm() / f.norm();
                temp = temp / relative_residual;
                std::cout << "iterations " << i << " Relative residual: " << relative_residual << " Residual reduction rate: " << temp << std::endl;
                std::cout << "Time elapsed: " << timer.elapsed() << " ms" << std::endl;
                if (relative_residual <= MG.relative_accuracy)
                {
                    break; // 达到收敛条件，退出迭代
                }
                temp = relative_residual;
            }
        }

        if (MG.type == MGParams_Two_Dimension::FMG)
        {
            double temp = (f - A * U).norm() / f.norm();
            double relative_residual = (f - A * U).norm() / f.norm();
            Timer timer;
            for (int i = 0; i < MG.max_iterations; i++)
            {
                U = FMG(A, f, 0);
                relative_residual = (f - A * U).norm() / f.norm();
                temp = temp / relative_residual;
                std::cout << "iterations " << i << " Relative residual: " << relative_residual << " Residual reduction rate: " << temp << std::endl;
                std::cout << "Time elapsed: " << timer.elapsed() << " ms" << std::endl;
                if (relative_residual <= MG.relative_accuracy)
                {
                    break; // 达到收敛条件，退出迭代
                }
                temp = relative_residual;
            }
        }
        return;
    }
    void Solve()
    {
        Eigen::SparseLU<Eigen::SparseMatrix<double>> solver;
        solver.compute(A);
        if (solver.info() != Eigen::Success)
        {
            std::cerr << "Error: Sparse matrix decomposition failed.\n";
            exit(1);
        }

        U = solver.solve(f);
        if (solver.info() != Eigen::Success)
        {
            std::cerr << "Error: Sparse matrix solver failed.\n";
            exit(1);
        }
    }

    const Eigen::VectorXd &GetSolution() const
    {
        return U;
    }
};

#endif // MULTIGRID_TWO_DIMENSION_HPP